Skip to content

perf(linear): share one eccentric-radius grid across an MGE basis (PyAutoArray#507) - #590

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/numba-hst-curvature-matrix-phase2
Aug 28, 2026
Merged

perf(linear): share one eccentric-radius grid across an MGE basis (PyAutoArray#507)#590
Jammy2211 merged 1 commit into
mainfrom
feature/numba-hst-curvature-matrix-phase2

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

The PyAutoGalaxy half of PyAutoLabs/PyAutoArray#507 phase 2. Step 0 of that
issue measured the MGE operated mapping matrix at 0.209 s of a 0.633 s HST
rectangular numba CPU likelihood evaluation, and — this is the part that was not
guessable — 78 % of it is profile evaluation, not the PSF convolution, which
phase 1's batching had already reduced to 0.054 s.

A multi-Gaussian expansion is a LightProfileLinearObjFuncList of tens of
Gaussian profiles which share a centre and ell_comps and differ only in
sigma. Building its mapping matrix evaluated image_2d_from once per profile,
and every one of those calls redid the same two grid calculations: the
reference-frame transform (an arctan2, a sin and a cos per coordinate) and
the eccentric radii formed from it. Both dwarf the profile itself, which is one
exp. cProfile on the HST cell: the transform alone was 951 us of the 1.3 ms
each profile took.

LightProfileLinearObjFuncList._image_slim_list_from now groups the profiles by
(class, centre, ell_comps) and computes each group's transform and radii once,
then evaluates every profile's image_2d_via_radii_from against its group's
radii. Both the data grid and the blurring grid go through it, so the numpy
batched-convolution path in operated_mapping_matrix_override picks it up too.

Grouping rather than an all-or-nothing check is deliberate. The workspace's
canonical MGE recipe (features/multi_gaussian_expansion/modeling.py) stacks two
sets of 30 Gaussians which share a centre but carry their own ell_comps, and
GalaxiesToInversion lands both sets in a single func list — an all-or-nothing
test would fall back to 60 independent evaluations on exactly the model users are
told to write.

Bit-identical, not approximate. The hoisted quantities are produced by
calling the same methods, with the same inputs, that each profile would have
called on itself.

Measured on the HST breakdown cell, 60 linear Gaussians, OMP_NUM_THREADS=1:
60x image_2d_from(grid) 0.0867 -> 0.0168 s, 60x
image_2d_from(blurring_grid) 0.0346 -> 0.0092 s. Paired in the harness, the
MGE step falls 0.194 -> 0.084 s at HST rectangular (2.3x) and 0.080 -> 0.032 s at
euclid (2.5x); with the PyAutoArray half it takes the whole HST evaluation
0.6214 -> 0.3334 s, which is the issue's goal.

API Changes

None — internal changes only. Three new private members on
LightProfileLinearObjFuncList (_shared_eccentric_radii_index_groups,
_image_slim_list_from, _shared_eccentric_radii_image_slim_list_from) and one
private module function (_gaussian_image_2d_from). mapping_matrix and
operated_mapping_matrix_override keep their signatures and return
bit-identical values.

See full details below.

Test Plan

  • pytest test_autogalaxy1149 passed (1144 before; +5).
  • test__mapping_matrix__shared_geometry_matches_per_profile_loop — an MGE
    basis matches the per-profile loop bit-identically (np.array_equal)
    on both the data grid and the blurring grid, and through
    operated_mapping_matrix_override against the per-profile
    psf.convolved_image_from.
  • ..._spherical_gaussians_share_radii_and_match_the_loop,
    ..._two_basis_mge_is_grouped_by_ell_comps (two groups plus a Sersic
    taking the loop), ..._mixed_geometry_falls_back (different centre /
    different ell_comps / different class),
    ..._non_gaussian_and_single_profile_take_the_loop.
  • Control runs recorded — each of these deliberately wrong variants fails
    the new tests: dropping the transform; using elliptical instead of
    eccentric radii; transforming the un-over-sampled grid; removing the
    image_2d_from-identity membership test.
  • Bit-identity on real data (HST, 60 Gaussians): the mapping matrix
    (60 x 15361), the blurring stack (60 x 5960) and the MGE operated mapping
    matrix all np.array_equal to the per-profile loop; the three pinned
    profiling log likelihoods unchanged to every recorded digit.

Why the membership test is image_2d_from identity

A group needs numpy, a Grid2D, two or more profiles, one class, an identical
centre and ell_comps, and — the load-bearing part — that the class's
image_2d_from is Gaussian.image_2d_from, tested by function-object
identity rather than isinstance. That is what makes "the transform and the
eccentric radii are the whole of the shared work" true:

  • subclasses which inherit it (linear, operated, spherical Gaussians) match;
  • GaussianMultipole, which overrides it to perturb the radius, does not;
  • Sersic and its children, which also share a geometry in a Sersic basis, do
    not — they already avoid the polar transform via
    _eccentric_radii_grid_from_cartesian, a branch taken only when the grid
    handed to them is not pre-transformed, so hoisting a transformed grid into
    them would be both slower and a different floating-point expression.

Follow-up worth filing: giving Gaussian that same Cartesian shortcut would
remove the polar transform for every Gaussian evaluation, not just an MGE
basis — but it moves values, so it is a separate, wider-blast-radius task.

Merge order

Merge PyAutoLabs/PyAutoArray#508 before this PR. This repo's CI clones
PyAutoArray from source and checks out the same-named branch when it exists, so
this PR is currently testing against the PyAutoArray branch; once #508 merges and
that branch is deleted, it falls back to PyAutoArray main. No test here depends
on #508's new OverSampler behaviour, so this PR is green either way — the order
is about not leaving the paired measurement half-landed. The workspace PR
PyAutoLabs/autolens_profiling#190 merges last, behind both.

Heart

Heart is RED at ship time for a pre-existing reason unrelated to this task.
Verbatim from pyauto-heart readiness --json (2026-08-28T21:34:42Z, score 45):

  • red_reasons: "release validation FAILED (stage integrate)"
  • yellow_reasons: "workspace validation not passing (2 failed, cloud#33179766004: autolens_test scripts/imaging/rectangular_mge.py, autolens_test scripts/imaging/rectangular_mge_rtu.py)",
    "manifest drift: session-start hooks (generated) — 32 mismatch(es) vs PyAutoMind/repos.yaml"

Human authorisation to ship over it, given 2026-08-28, verbatim:

prm and then kick off phase 2, I authorize the heart RED thing

Full API Changes (for automation & release notes)

Added

  • autogalaxy.profiles.light.linear.abstract._gaussian_image_2d_from() — private; returns the Gaussian.image_2d_from function object used as the group-membership token.
  • LightProfileLinearObjFuncList._shared_eccentric_radii_index_groups — private cached property; the profile indices grouped by shared geometry.
  • LightProfileLinearObjFuncList._image_slim_list_from(grid, xp=np) — private; the slim image of every profile, grouped where possible.
  • LightProfileLinearObjFuncList._shared_eccentric_radii_image_slim_list_from(grid, light_profile_list) — private; one group's images from one shared transform and radius grid.

Changed Behaviour

  • LightProfileLinearObjFuncList.mapping_matrix and the numpy branch of operated_mapping_matrix_override route their per-profile loops through _image_slim_list_from. Returned values are bit-identical.

Migration

  • None required.

Generated by the PyAutoLabs agent workflow.

…AutoArray#507)

A multi-Gaussian expansion is a `LightProfileLinearObjFuncList` of tens of
`Gaussian` profiles which share a `centre` and `ell_comps` and differ only in
`sigma`. Building its mapping matrix evaluated `image_2d_from` once per
profile, and every one of those calls redid the same two grid calculations:

  1. the reference-frame transform (`transform_grid_2d_to_reference_frame`),
     which costs an `arctan2`, a `sin` and a `cos` per coordinate;
  2. the eccentric radii formed from it.

Both dwarf the profile itself, which is one `exp`. Measured on the HST
breakdown cell (60 Gaussians, 17980 over-sampled coordinates, cProfile,
OMP_NUM_THREADS=1) the transform alone was 951 us of the 1.3 ms each profile
took, i.e. ~70 % of the MGE's profile-evaluation cost.

`LightProfileLinearObjFuncList._image_slim_list_from` now groups the profiles
by the geometry they share, computes each group's transform and radii once, and
evaluates every profile's `image_2d_via_radii_from` against its group's radii.
Grouping rather than an all-or-nothing check is deliberate: the workspace's
canonical MGE recipe stacks two sets of 30 Gaussians which share a centre but
carry their own `ell_comps`, and `GalaxiesToInversion` lands both sets in one
func list -- an all-or-nothing test would fall back to 60 independent
evaluations on exactly the model users are told to write.

Binning stays a per-profile call (`OverSampler.binned_array_2d_from` takes a 1D
array), and so does the `exp`: that arithmetic is transcendental-bound rather
than overhead-bound, so a single (over-sampled pixels, profiles) block buys
nothing and would have to be sliced column-wise to bin. Both the data grid and
the blurring grid go through it, so the numpy batched-convolution path in
`operated_mapping_matrix_override` picks it up as well.

Bit-identical, not approximate: the hoisted quantities are produced by calling
the same methods, with the same inputs, that each profile would have called on
itself. Asserted with `np.array_equal` on the HST MGE (60 x 15361 and
60 x 5960) and in the unit tests; the three pinned harness log likelihoods are
unchanged to every recorded digit.

The membership test is deliberately narrow. A group needs numpy, a `Grid2D`,
two or more profiles, one class, an identical `centre` and `ell_comps`, and --
the load-bearing part -- that the class's `image_2d_from` *is*
`Gaussian.image_2d_from`, tested by function-object identity rather than
`isinstance`. That is what makes "the transform and the eccentric radii are the
whole of the shared work" true:

  - subclasses which inherit it (linear, operated, spherical Gaussians) match;
  - `GaussianMultipole`, which overrides it to perturb the radius, does not;
  - `Sersic` and its children, which also share a geometry in a Sersic basis,
    do not -- they already avoid the polar transform via
    `_eccentric_radii_grid_from_cartesian`, a branch taken only when the grid
    handed to them is *not* pre-transformed, so hoisting into them would be
    both slower and a different floating-point expression. (Giving `Gaussian`
    that same Cartesian shortcut would be a larger, value-changing lever and is
    left as a follow-up.)

Everything outside a group is evaluated by the per-profile loop, unchanged.

Measured on the HST breakdown cell, 60 linear Gaussians, OMP_NUM_THREADS=1:
60x `image_2d_from(grid)` 0.0867 -> 0.0168 s and 60x
`image_2d_from(blurring_grid)` 0.0346 -> 0.0092 s. In the paired harness runs
the MGE operated-mapping-matrix step falls 0.194 -> 0.084 s at HST rectangular
(2.3x) and 0.080 -> 0.032 s at euclid (2.5x).

Tests (`test_autogalaxy/profiles/light/linear/test_abstract.py`): shared
geometry matches the per-profile loop bit-identically on both grids and through
`operated_mapping_matrix_override`; the spherical Gaussian twin; a two-basis
MGE with a Sersic mixed in is grouped by `ell_comps` and still matches; mixed
geometry (different centre / different ell_comps / different class) falls back
and matches; `Sersic`, `GaussianMultipole` and a one-profile list take the
loop. Control runs recorded: dropping the transform, using elliptical instead
of eccentric radii, transforming the un-over-sampled grid, and removing the
`image_2d_from`-identity guard each fail these tests.

`pytest test_autogalaxy`: 1149 passed (1144 before, +5).

Refs PyAutoLabs/PyAutoArray#507

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

Upstream: PyAutoLabs/PyAutoArray#508. Workspace: PyAutoLabs/autolens_profiling#190. Merge order: PyAutoArray#508 -> this PR -> autolens_profiling#190.

@Jammy2211
Jammy2211 merged commit 9031712 into main Aug 28, 2026
4 checks passed
@Jammy2211
Jammy2211 deleted the feature/numba-hst-curvature-matrix-phase2 branch August 28, 2026 22:52
@Jammy2211 Jammy2211 removed the pending-release PR queued for the next release build label Sep 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant